package com.educationagent.service.impl;

import com.educationagent.service.AgentService;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.DocumentSplitter;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.loader.UrlDocumentLoader;
import dev.langchain4j.data.document.parser.apache.pdfbox.ApachePdfBoxDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.HuggingFaceTokenizer;
import dev.langchain4j.store.embedding.EmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.concurrent.CompletableFuture;

@Service
public class AgentServiceImpl implements AgentService {

    @Autowired
    private EmbeddingModel embeddingModel;

    @Autowired
    private EmbeddingStore embeddingStore;

     public Boolean load_document(String fileURL) {

        try {
            Document document = UrlDocumentLoader.load(fileURL, new ApachePdfBoxDocumentParser());


            // 递归分割
            DocumentSplitter splitter = DocumentSplitters.recursive(300, 30,new HuggingFaceTokenizer());
            //段落分割 使用HuggingFaceTokenizer:token分词器
       /* DocumentByParagraphSplitter splitter = new DocumentByParagraphSplitter(
                300, 30, new HuggingFaceTokenizer()
        );*/


            //文档分割
            List<TextSegment> segments = splitter.split(document);
            //文档向量化
            List<Embedding> embeddings = embeddingModel.embedAll(segments).content();

            embeddingStore.addAll(embeddings,segments);
        }catch (Exception e) {
            return false;
        }
        return true;

    }


    @Async
    public CompletableFuture<Boolean> longRunningMethod(String fileURL) throws InterruptedException {
        // 执行异步任务
        Boolean flag = load_document(fileURL);
        return CompletableFuture.completedFuture(flag);
    }


}
